import pandas as pd
import numpy as np
import torch
from get_result import angle, fa
import cv2
import os
import torch
import time
# from progressbar import ProgressBar
from data_aug import resize, center_crop, random_crop, aug_predict, multi_model
from models.model import model_resnet_unet_GN, model_hrnet, model_Unet_efficient_b3
from make_heatmap import heatmap_to_anno


def compute_loss(pred, label, original_size=None, resize_size=None, distance=1):
    # 传入pred和lable的单张图片的坐标
    ## 用来记录距离小于指定数值的正样本个数
    positive = 0
    ## 用来记录l2 loss
    l2_loss = []
    assert len(pred) == len(label), 'the len of pred and label shoule be same'
    ## 输入全部的pred和label，用list中包含tuple的形式，其中像素间距为0.139mm
    for index in range(len(pred)):
        single_pred = pred[index]
        single_label = label[index]
        ## 遇到缺失的label不进行计算
        if 0 in single_label:
            continue
        ## 计算两个点之间的像素距离。注意根据原始尺寸和resize后的尺寸对坐标值进行转换。
        if original_size is not None:
            loss = (((single_pred[0] - single_label[0]) / resize_size[0] * original_size[0]) ** 2 +
                    ((single_pred[1] - single_label[1]) / resize_size[1] * original_size[1]) ** 2) ** 0.5
        else:
            ## 如果使用原图上的坐标。直接进行计算即可，不需要进行转换坐标。
            loss = (((single_pred[0] - single_label[0])) ** 2 +
                    ((single_pred[1] - single_label[1])) ** 2) ** 0.5
        ## 将计算出的像素距离转换成物理距离
        distance_loss = float(loss) * 0.143
        # print('distance loss',distance_loss)
        l2_loss.append(distance_loss)
        if distance_loss <= distance:
            positive += 1
    return positive, l2_loss


def predict(model_1,model_2, csv_file_dir, image_size_1,image_size_2, show_pred=False,
            test_aug=False, print_l2_loss=False, Use_Cewei=False, model2=None):
    csv_data = pd.read_csv(csv_file_dir).fillna('[0,0,0]')
    name_list = (csv_data.columns.values.tolist())[2:]
    pred_csv = pd.DataFrame(columns=['image_path'] + name_list)
    with torch.no_grad():
        model_1.eval()
        model_2.eval()
        # progressbar = ProgressBar()
        for index in range(len(csv_data)):
            image_dir = csv_data.iloc[index].loc['image_path']
            image_1 = cv2.imread(image_dir)
            original_image = image_1.copy()
            image_2 = cv2.imread(image_dir)
            # opencv读取数据为BRG，转成RGB
            image_1 = cv2.cvtColor(image_1, cv2.COLOR_BGR2RGB)
            image_2 = cv2.cvtColor(image_2, cv2.COLOR_BGR2RGB)
            original_size = image_1.shape[:-1]
            ## 读取label
            label = []
            for i in csv_data.iloc[index, 2:]:
                label.append(
                    (int(float(i.split(',')[1].replace(']', ''))), int(float(i.split(',')[0].replace('[', '')))))
            
            image_1, label_1 = resize(image_1, label, image_size_1)
            image_2, label_2 = resize(image_2, label, image_size_2)

            
            input_1 = torch.from_numpy(np.transpose(image_1, axes=(2, 0, 1))).float().cuda()
            input_1 = input_1[np.newaxis, ...]

            input_2 = torch.from_numpy(np.transpose(image_2, axes=(2, 0, 1))).float().cuda()
            input_2 = input_2[np.newaxis, ...]
            

            pred_1 = model_1(input_1)
            pred_2 = model_2(input_2)
            
            ## 注意pred[0]去掉batch维度
            pred_1 = heatmap_to_anno(pred_1[0].cpu())
            original_pred_1 = [(int(float(p[0] / image_size_1[0] * original_size[0])),
                              int(float(p[1] / image_size_1[1] * original_size[1]))) for p in pred_1]

            ## 注意pred[0]去掉batch维度
            pred_2 = heatmap_to_anno(pred_2[0].cpu())
            original_pred_2 = [(int(float(p[0] / image_size_2[0] * original_size[0])),
                                int(float(p[1] / image_size_2[1] * original_size[1]))) for p in pred_2]

            original_pred = original_pred_1 + original_pred_2
            pred_dict = dict(zip(name_list,original_pred))
            if show_pred:
                for index in range(len(original_pred)):
                    single_pred = original_pred[index]
                    single_label = label[index]
                    single_label = (int(single_label[0]), int(single_label[1]))
                    if 0 in single_label:
                        continue
                    ##  注意对坐标进行反转成坐标系中的方式，来进行画点
                    cv2.circle(original_image, single_label[::-1], 1, (255, 128, 255), thickness=9)
                    cv2.circle(original_image, single_pred[::-1], 1, (128, 255, 0), thickness=9)
                    
                    cv2.putText(original_image, name_list[index], single_label[::-1], cv2.FONT_HERSHEY_PLAIN, 3, (255, 128, 255),
                                1)
                    # cv2.putText(image, name_list[index], tuple(map(lambda x:x+15,single_pred[::-1])), cv2.FONT_HERSHEY_PLAIN, 1, (128, 255, 0), 1)
                # original_image = cv2.resize(original_image,(512,768))
                # cv2.imshow('original_image', original_image)
                # cv2.waitKey(0)
            cv2.imwrite(os.path.join('../data/process_data/show/pred', image_dir.split('/')[-1]), original_image)
            pred_dict['image_path'] = image_dir
            pred_dict = {k:str(v) for k,v in pred_dict.items()}
            pred_csv = pred_csv.append(pd.DataFrame(pred_dict,index=[0]))
    pred_csv.to_csv('./pred_csv.csv',index=False)



if __name__ == '__main__':
    # for i in range(31,249,2):
    #     i = '199'
    
    start_time = time.time()
    
    Image_Size_2 = (768,384)
    Image_Size_1 = (384, 256)
    trained_model_dir_1 = '../model_save/baseline_1_0618/549.pth'
    trained_model_dir_2 = '../model_save/baseline_2_0618/83.pth'
    
    # val_csv_dir = '../data/csv_file/val_data.csv'
    Image_Path = '../data/original_data'
    # trained_model_dir = '../model_save/0610_cut/139.pth'
    val_csv_dir = '../data/process_data/label_file/test.csv'
    # val_csv_dir = '../data/label_file/0522/cewei/val_cewei_cut_heatmap.csv'
    # model = model_cpn50(image_size = Image_Size)
    # model = model_simple_pose_res101()
    model_1 = model_resnet_unet_GN(layer='resnet34', use_aspp=True, num_classes=8, pre_train=False)
    model_2 = model_resnet_unet_GN(layer='resnet34', use_aspp=True, num_classes=14, pre_train=False)

    model_1.load_state_dict(torch.load(trained_model_dir_1))
    model_2.load_state_dict(torch.load(trained_model_dir_2))

    predict(model_1,model_2, val_csv_dir, Image_Size_1,Image_Size_2,
                                         show_pred=True, test_aug=False, print_l2_loss=False,
                                         Use_Cewei=False, model2=None)
    # print('use time ---',time.time() - start_time)